import os
import datetime as dt
import time
import pickle
import pandas as pd
import numpy as np
import math
from tqdm import tqdm
from sklearn import svm, preprocessing, metrics
from datahelper import StockDataHelper

class TrainModel:
    def __init__(self,window_size,next_n_day):
        # 根据过往𝑘天的数据，来预测股票𝑛天后的涨跌情况
        self.window_size = window_size  # k
        self.next_n_day = next_n_day  # n
        # 划分训练集测试集的年份
        self.split_year = 2013

        datahelper = StockDataHelper()
        file_path = os.path.join(os.getcwd(),'..', 'data', 'data.pkl')
        datahelper.init_data(file_path)
        self.data = datahelper.data_by_date
        self.trade_day_list = datahelper.trade_day_list

        self.train_data, self.test_data = self.load_train_and_test_data()

    def load_train_and_test_data(self):
        data = self.data

        train_index, test_index = [], []
        for date in data.keys():
            if int(date[:4]) <= self.split_year:
                train_index.append(date)
            else:
                test_index.append(date)

        train_index = sorted(train_index)
        test_index = sorted(test_index)
        train_index.sort()
        test_index.sort()

        train_data, test_data = {}, {}
        for date in train_index:
            train_data[date] = data[date]
        for date in test_index:
            test_data[date] = data[date]

        return train_data, test_data

    def is_valid(self, data):
        for n in data:
            if n is None or n != n:
                return False
        return True

    def reform_dataframe(self,trade_day):
        """
        :param dataframe:一个trade day的dataframe
        :return:
        """
        data = self.data
        trade_day_list = self.trade_day_list
        window_size = self.window_size
        next_n_day = self.next_n_day

        dataframe = data[trade_day]

        trade_day_idx = trade_day_list.index(trade_day)
        factor_dict = {}
        label_dict = {}
        for i in range(len(dataframe)):
            stock = dataframe.index[i]
            factor = []
            for j in range(window_size):  # 过去window_size天的所有数据
                idx = trade_day_idx - window_size + j  # 包含当天数据
                pre_day = trade_day_list[idx]
                pre_df = data[pre_day]

                for keyword in ['volume', 'money', 'WVAD', 'OBV']:
                    pre_df[keyword] = pre_df[keyword] / 1000000

                pre_factor = pre_df.loc[stock].tolist()
                factor += pre_factor

            # label
            current_close = dataframe.loc[stock, 'close']
            next_n_day_close = data[trade_day_list[trade_day_idx + next_n_day]].loc[stock, 'close']
            label = 1 if next_n_day_close > current_close else 0

            if not self.is_valid(factor + [current_close, next_n_day_close]):  # 有空值则跳过
                continue
            factor_dict[stock] = factor
            label_dict[stock] = label

        return factor_dict, label_dict

    def preprocess(self,data):
        """
        数据预处理
        :param data: dict{date : dataframe,...}
        :return:
        """
        trade_day_list = self.trade_day_list
        window_size = self.window_size
        next_n_day = self.next_n_day

        x_data, y_data = None, None

        data_day_list = list(data.keys())
        data_day_list.sort()
        for i in tqdm(range(len(data_day_list))):
            if i < window_size or i >= (len(data_day_list) - next_n_day):  # 防止下标越界
                continue
            trade_day = data_day_list[i]
            factor_dict, label_dict = self.reform_dataframe(trade_day)
            for stock in factor_dict.keys():
                if factor_dict[stock]:
                    print(factor_dict[stock])
                    print(label_dict[stock])
                    x = np.array(factor_dict[stock], dtype=float).reshape(-1, len(factor_dict[stock]))
                    y = np.array([label_dict[stock]], dtype=float)

                if x_data is None:
                    x_data, y_data = x, y
                else:
                    x_data = np.concatenate((x_data, x), axis=0)
                    y_data = np.concatenate((y_data, y), axis=0)

        return x_data, y_data

    def train_model(self):
        train_data = self.train_data
        model = svm.SVC(kernel="rbf", C=100.0)
        train_x, train_y = self.preprocess(train_data)
        model.fit(train_x, train_y)
        return model

    def evaluate(self, test_y, pred_y):
        # count accuracy
        acc_num = 0
        for i in range(len(test_y)):
            acc_num += (1 if pred_y[i] == test_y[i] else 0)
        acc_rate = acc_num / len(test_y)

        # count precision, recall, f1-score
        def prf_measure(golden_label, pred_label, mark):
            TP, FP, TN, FN = 0, 0, 0, 0

            for i in range(len(golden_label)):
                gold, pred = golden_label[i], pred_label[i]
                if gold == mark and pred == mark:
                    TP += 1
                elif gold == mark and pred != mark:
                    FN += 1
                elif gold != mark and pred == mark:
                    FP += 1
                elif gold != mark and pred != mark:
                    TN += 1

            prec, recall = TP / max((TP + FP), 1), TP / max((TP + FN), 1)
            if prec * recall == 0:
                f1 = 0
            else:
                f1 = 2 * (prec * recall) / (prec + recall)

            return prec, recall, f1

        prec, recall, f1 = prf_measure(test_y, pred_y, mark=1)

        print("Accuracy rate: %.2f%%" % (acc_rate * 100))
        print("Precision: %.2f   Recall: %.2f   F1: %.2f" % (prec * 100, recall * 100, f1 * 100))

    def predict(self,model):
        test_data = self.test_data
        test_x, test_y = self.preprocess(test_data)
        pred_y = model.predict(test_x)
        return test_y, pred_y

    def save_model(self, model):
        model_name = "svm_win%d_next%d.md" % (self.window_size, self.next_n_day)
        model_path = os.path.join(os.getcwd(),"strategy","svm",model_name)
        if os.path.exists(model_path):
            os.remove(model_path)
        with open(model_path, 'wb') as f:
            pickle.dump(model, f)

    def run(self):
        model = self.train_model()
        test_y, pred_y = self.predict(model)
        self.evaluate(test_y, pred_y)
        self.save_model(model)

if __name__ == '__main__':
    trainmodel = TrainModel(2, 3)
    trainmodel.run()